#include <llvm/IR/Function.h>
#include <llvm/IR/IRBuilder.h>
#include <llvm/IR/PassManager.h>
#include <llvm/Passes/PassBuilder.h>
#include <llvm/Support/raw_ostream.h>

using namespace llvm;

namespace {

class SimpleFunctionPass : public PassInfoMixin<SimpleFunctionPass> {
public:
  PreservedAnalyses run(Function &F, FunctionAnalysisManager &) {
    unsigned replacedCount = 0;

    for (auto &BB : F) {
      for (auto instIt = BB.begin(), instEnd = BB.end(); instIt != instEnd;) {
        Instruction &I = *instIt++;

        // 查找形如 `add x, x` 的指令，并替换为 `shl x, 1`
        if (auto *binOp = dyn_cast<BinaryOperator>(&I)) {
          if (binOp->getOpcode() == Instruction::Add &&
              binOp->getOperand(0) == binOp->getOperand(1) &&
              binOp->getType()->isIntegerTy()) {
            IRBuilder<> builder(binOp);
            Value *shifted =
                builder.CreateShl(binOp->getOperand(0),
                                  ConstantInt::get(binOp->getType(), 1),
                                  "double.by.shift");
            binOp->replaceAllUsesWith(shifted);
            binOp->eraseFromParent();
            ++replacedCount;
          }
        }
      }
    }

    if (replacedCount > 0) {
      errs() << "[SimpleFunctionPass] Replaced " << replacedCount
             << " additions with shifts in function " << F.getName() << "\n";
      return PreservedAnalyses::none();
    }

    errs() << "[SimpleFunctionPass] No opportunities in " << F.getName()
           << "\n";
    return PreservedAnalyses::all();
  }
};

} // end anonymous namespace

// 向新 Pass 管理器注册
llvm::PassPluginLibraryInfo getSimpleFunctionPassPluginInfo() {
  return {
      LLVM_PLUGIN_API_VERSION, "SimpleFunctionPass", LLVM_VERSION_STRING,
      [](PassBuilder &PB) {
        PB.registerPipelineParsingCallback(
            [](StringRef Name, FunctionPassManager &FPM,
               ArrayRef<PassBuilder::PipelineElement>) {
              if (Name == "simple-func-pass") {
                FPM.addPass(SimpleFunctionPass());
                return true;
              }
              return false;
            });
      }};
}

extern "C" LLVM_ATTRIBUTE_WEAK ::llvm::PassPluginLibraryInfo
llvmGetPassPluginInfo() {
  return getSimpleFunctionPassPluginInfo();
}
